1366D - Two Divisors - CodeForces Solution

#include <iostream>
#include <utility>
#include <vector>
#include <algorithm>
#define DEBUG 0
using namespace std;

const int MAX = 10'000'000;
const pair<int, int> INVALID_PAIR(-1, -1);

int largest_prime_fact[MAX+1] = {0};

int gcd(int a, int b) {
	while (b > 0) {
		int r = a % b;
		a = b;
		b = r;
	}
	return a;
}

void init_largest_prime_fact_arr() {
	for (int i = 2; i <= MAX; i++) {
		if (largest_prime_fact[i] == 0) {
			for (int j = i; j <= MAX; j += i) {
				largest_prime_fact[j] = i;
			}
		}
	}
}

pair<int, int> get_two_divisors(int n) {
	// return (d1, d2) where d1, d2 are divisors of n and gcd(d1 + d2, n) == 1
	vector<pair<int, int>> n_prime_facts;
	for (int n_ = n, p; n_ > 1; n_ /= p) {
		p = largest_prime_fact[n_];
		if (n_prime_facts.empty() || n_prime_facts.back().first != p) {
			n_prime_facts.push_back({p, 1});
		} else {
			n_prime_facts.back().second++;
		}
	}
	if (n_prime_facts.size() == 1) {
		// n = p**k (p is prime, k > 0)
		return INVALID_PAIR;
	}
	reverse(n_prime_facts.begin(), n_prime_facts.end());
	if (DEBUG) {
		cout << n << " ->";
		for (auto p: n_prime_facts) cout << " " << p.first << "**" << p.second;
		cout << endl;
	}
	for (int i = 0, d1; i < n_prime_facts.size() && (d1 = n_prime_facts[i].first); i++) {
		for (int d1_pwr = 1, k1 = 0; d1_pwr *= d1, ++k1 <= n_prime_facts[i].second;) {
			if (gcd(d1_pwr + n / d1_pwr, n) == 1) {
				return {d1_pwr, n / d1_pwr};
			}
			for (int j = i + 1, d2; j < n_prime_facts.size() && (d2 = n_prime_facts[j].first); j++) {
				for (int d2_pwr = 1, k2 = 0; d2_pwr *= d2, ++k2 <= n_prime_facts[j].second;) {
					if (gcd(d1_pwr + d2_pwr, n) == 1) {
						return {d1_pwr, d2_pwr};
					}
					if (gcd(n / d1_pwr + d2_pwr, n) == 1) {
						return {n / d1_pwr, d2_pwr};
					}
					if (gcd(d1_pwr + n / d2_pwr, n) == 1) {
						return {d1_pwr, n / d2_pwr};
					}
					if (gcd(n / d1_pwr + n / d2_pwr, n) == 1) {
						return {n / d1_pwr, n / d2_pwr};
					}
				}
			}
		}
	}
	return INVALID_PAIR;
}

int main() {
	init_largest_prime_fact_arr();
	int t, n;
	vector<int> d1_arr, d2_arr;
	for (cin >> t; t-- && cin >> n; ) {
		auto d_pair = get_two_divisors(n);
		d1_arr.push_back(d_pair.first);
		d2_arr.push_back(d_pair.second);
		if (DEBUG && t >= 500000 - 1 - 143 && d_pair == INVALID_PAIR) {
			cout << n << " " << d_pair.first << " " << d_pair.second << endl;
		}
	}
	for (int d1: d1_arr) cout << d1 << " "; cout << endl;
	for (int d2: d2_arr) cout << d2 << " "; cout << endl;
}